﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Runtime.InteropServices;
using System.Web;
using Ast.Core.Services.Extensions;
using Ast.Core.Services.ObjectDefinitions;
using Ast.Framework;
using Ast.Framework.Caching;
using Ast.Framework.Services;
using Castle.Core.Logging;

namespace Ast.Core.Services
{
    public interface IObjectTypeService : IDependency
    {
        IEnumerable<ObjectTypeData> ListObjectTypes();
        IList<ObjectTypeData> ListObjectTypesInCache(bool reloadCache);

        void ListObjectsInDirectory(IntPtr dirPtr, Func<ObjectDirectory, ObjectDirectoryEntry, bool> objProcFunc);

        ObjectData<ObjectHeader> GetObjectHeader(IntPtr objPtr);
        T? GetSpecificObjectHeaderInfo<T>(IntPtr objPtr) where T : struct, ISpecificObjectHeaderInfo;
        ObjectTypeData GetObjectType(IntPtr objPtr);

        string GetObjectName(IntPtr objPtr);
        string GetFileObjectName(IntPtr objPtr);
        string GetKeyObjectName(IntPtr objPtr);
        string GetSymbolicLinkObjectLinkTarget(IntPtr objPtr);
    }

    public class ObjectTypeService : IObjectTypeService
    {
        private readonly IDeviceController _deviceController;
        private readonly IKernelAddressFinder _kernelAddressFinder;
        private readonly ICacheManager _cacheManager;
        private readonly ISignals _signals;
        private readonly IClock _clock;

        private static readonly int[] ObpInfoMaskToOffset = new int[128];

        static ObjectTypeService()
        {
            InitObpInfoMaskToOffset();
        }

        private static void InitObpInfoMaskToOffset()
        {
            for (InfoMaskType i = 0; (uint) i < ObpInfoMaskToOffset.Length; i++)
            {
                var offset = 0;
                if ((i & InfoMaskType.CreatorInfo) == InfoMaskType.CreatorInfo)
                {
                    offset += Marshal.SizeOf<ObjectHeaderCreatorInfo>();
                }
                if ((i & InfoMaskType.NameInfo) == InfoMaskType.NameInfo)
                {
                    offset += Marshal.SizeOf<ObjectHeaderNameInfo>();
                }
                if ((i & InfoMaskType.HandleInfo) == InfoMaskType.HandleInfo)
                {
                    offset += Marshal.SizeOf<ObjectHeaderHandleInfo>();
                }
                if ((i & InfoMaskType.QuotaInfo) == InfoMaskType.QuotaInfo)
                {
                    offset += Marshal.SizeOf<ObjectHeaderQuotaInfo>();
                }
                if ((i & InfoMaskType.ProcessInfo) == InfoMaskType.ProcessInfo)
                {
                    offset += Marshal.SizeOf<ObjectHeaderProcessInfo>();
                }
                if ((i & InfoMaskType.AuditInfo) == InfoMaskType.AuditInfo)
                {
                    offset += Marshal.SizeOf<ObjectHeaderAuditInfo>();
                }
                ObpInfoMaskToOffset[(uint) i] = offset;
            }
        }

        public ObjectTypeService(IDeviceController deviceController, IKernelAddressFinder kernelAddressFinder, ICacheManager cacheManager, ISignals signals, IClock clock)
        {
            _deviceController = deviceController;
            _kernelAddressFinder = kernelAddressFinder;
            _cacheManager = cacheManager;
            _signals = signals;
            _clock = clock;

            Logger = NullLogger.Instance;
        }

        public ILogger Logger { get; set; }

        public IEnumerable<ObjectTypeData> ListObjectTypes()
        {
            var obpTypeObjectType = _kernelAddressFinder.FindObpTypeObjectType();
            var objectType = _deviceController.ReadAsIntPtr(obpTypeObjectType).ReadAs<ObjectType>();
            var listHeader = objectType.TypeList.Flink;
            var currentAddress = listHeader;
            var currentObjectHeaderCreatorInfo = _deviceController.ReadAs<ObjectHeaderCreatorInfo>(currentAddress);
            while (currentObjectHeaderCreatorInfo.TypeList.Flink != listHeader)
            {
                var currentObjectHeaderAddress = currentAddress + Marshal.SizeOf<ObjectHeaderCreatorInfo>();
                var currrentObjectTypeAddress = currentObjectHeaderAddress + Marshal.SizeOf<ObjectHeader>();
                var currentObjectHeader = _deviceController.ReadAs<ObjectHeader>(currentObjectHeaderAddress);
                var currentObjectType = _deviceController.ReadAs<ObjectType>(currrentObjectTypeAddress);
                var name = _deviceController.ReadAsUnicodeString(currentObjectType.Name);

                Logger.InfoFormat("found type {0}, index 0x{1:x}, at 0x{2:x}", name, currentObjectType.Index, currentAddress.ToInt32());

                yield return new ObjectTypeData
                {
                    Name = name,
                    ObjectHeaderCreatorInfo = new ObjectData<ObjectHeaderCreatorInfo>(currentAddress, currentObjectHeaderCreatorInfo),
                    ObjectHeader = new ObjectData<ObjectHeader>(currentObjectHeaderAddress, currentObjectHeader),
                    ObjectType = new ObjectData<ObjectType>(currrentObjectTypeAddress, currentObjectType)
                };

                // next
                currentAddress = currentObjectHeaderCreatorInfo.TypeList.Flink;
                currentObjectHeaderCreatorInfo = _deviceController.ReadAs<ObjectHeaderCreatorInfo>(currentAddress);
            }
        }

        public IList<ObjectTypeData> ListObjectTypesInCache(bool reloadCache)
        {
            const string key = "ObjectTypeService.ListObjectTypesInCache";

            if (reloadCache) _signals.Trigger(key); // reload cache

            return _cacheManager.Get(key, context =>
            {
                context.Monitor(_signals.When(key)); // when reload required

                return ListObjectTypes().ToList();
            });
        }

        public void ListObjectsInDirectory(IntPtr dirPtr, Func<ObjectDirectory, ObjectDirectoryEntry, bool> objProcFunc)
        {
            var objectDirectory = _deviceController.ReadAs<ObjectDirectory>(dirPtr);
            foreach (var hashBucket in objectDirectory.HashBuckets.Where(hashBucket => hashBucket != IntPtr.Zero))
            {
                var goOn = true;

                var nextObjectDirectoryEntryPtr = hashBucket;
                do
                {
                    var objectDirectoryEntry = _deviceController.ReadAs<ObjectDirectoryEntry>(nextObjectDirectoryEntryPtr);
                    if (objectDirectoryEntry.Object == IntPtr.Zero) continue;

                    goOn = objProcFunc(objectDirectory, objectDirectoryEntry);
                    if (!goOn) break;

                    // next
                    nextObjectDirectoryEntryPtr = objectDirectoryEntry.ChainLink;
                } while (nextObjectDirectoryEntryPtr != IntPtr.Zero);

                if (!goOn) break;
            }
        }

        public ObjectData<ObjectHeader> GetObjectHeader(IntPtr objPtr)
        {
            var objectHeaderPtr = objPtr - Marshal.SizeOf<ObjectHeader>();
            var objectHeader = _deviceController.ReadAs<ObjectHeader>(objectHeaderPtr);
            return new ObjectData<ObjectHeader>(objectHeaderPtr, objectHeader);
        }

        public T? GetSpecificObjectHeaderInfo<T>(IntPtr objPtr) where T : struct, ISpecificObjectHeaderInfo
        {
            var objectHeaderData = GetObjectHeader(objPtr);
            var infoMask = (InfoMaskType) objectHeaderData.Object.InfoMask;

            var targetHeader = new T();
            var headerMask = targetHeader.Type;

            if ((infoMask & headerMask) == 0) return null;

            var maxMask = (byte) headerMask*2 - 1;
            var targetHeaderPtr = objectHeaderData.Address - ObpInfoMaskToOffset[(byte) infoMask & maxMask];
            return _deviceController.ReadAs<T>(targetHeaderPtr);
        }

        public ObjectTypeData GetObjectType(IntPtr objPtr)
        {
            var objectHeaderData = GetObjectHeader(objPtr);
            var objectTypes = ListObjectTypesInCache(false);
            return objectTypes.FirstOrDefault(x => x.ObjectType.Object.Index == objectHeaderData.Object.TypeIndex);
        }

        public string GetObjectName(IntPtr objPtr)
        {
            return _cacheManager.Get(string.Format("ObjectTypeService.GetObjectName.0x{0:x}", objPtr.ToInt32()), context =>
            {
                context.Monitor(_clock.When(TimeSpan.FromSeconds(5))); // cache for 5 secs

                var objectHeaderNameInfo = GetSpecificObjectHeaderInfo<ObjectHeaderNameInfo>(objPtr);
                if (objectHeaderNameInfo.HasValue)
                {
                    var objectName = _deviceController.ReadAsUnicodeString(objectHeaderNameInfo.Value.Name);
                    if (objectHeaderNameInfo.Value.Directory == IntPtr.Zero) return objectName;

                    var dirName = GetObjectName(objectHeaderNameInfo.Value.Directory);
                    if (dirName.Equals("\\")) dirName = string.Empty;
                    return string.Format("{0}\\{1}", dirName, objectName);
                }
                return string.Empty;
            });
        }

        public string GetFileObjectName(IntPtr objPtr)
        {
            return _cacheManager.Get(string.Format("ObjectTypeService.GetFileObjectName.0x{0:x}", objPtr.ToInt32()), context =>
            {
                context.Monitor(_clock.When(TimeSpan.FromSeconds(5))); // cache for 5 secs

                var fileObject = _deviceController.ReadAs<FileObject>(objPtr);
                var fileName = _deviceController.ReadAsUnicodeString(fileObject.FileName);
                if (fileObject.DeviceObject == IntPtr.Zero) return fileName;

                var deviceName = GetObjectName(fileObject.DeviceObject);
                return deviceName + fileName;
            });
        }

        public string GetKeyObjectName(IntPtr objPtr)
        {
            var keyBody = _deviceController.ReadAs<CmKeyBody>(objPtr);
            return ReadKeyName(keyBody.KeyControlBlock);
        }

        public string GetSymbolicLinkObjectLinkTarget(IntPtr objPtr)
        {
            var objectSymbolicLink = _deviceController.ReadAs<ObjectSymbolicLink>(objPtr);
            return _deviceController.ReadAsUnicodeString(objectSymbolicLink.LinkTarget);
        }

        private string ReadKeyName(IntPtr keyControlBlockPtr)
        {
            return _cacheManager.Get(string.Format("ObjectTypeService.ReadKeyName.0x{0:x}", keyControlBlockPtr.ToInt32()), context =>
            {
                context.Monitor(_clock.When(TimeSpan.FromSeconds(5))); // cache for 5 secs

                if (keyControlBlockPtr == IntPtr.Zero) return string.Empty;

                var keyControlBlock = _deviceController.ReadAs<CmKeyControlBlock>(keyControlBlockPtr);
                var nameBlock = _deviceController.ReadAs<CmNameControlBlock>(keyControlBlock.NameBlock);
                var name = _deviceController.ReadAsStringAnsi(keyControlBlock.NameBlock + Marshal.SizeOf<CmNameControlBlock>(), nameBlock.NameLength);

                return keyControlBlock.ParentKcb == IntPtr.Zero
                    ? string.Format("\\{0}", name)
                    : string.Format("{0}\\{1}", ReadKeyName(keyControlBlock.ParentKcb), name);
            });
        }
    }

    [Flags]
    public enum InfoMaskType : byte
    {
        CreatorInfo = 0x1,
        NameInfo = 0x2,
        HandleInfo = 0x4,
        QuotaInfo = 0x8,
        ProcessInfo = 0x10,
        AuditInfo = 0x40,
    }

    public class ObjectNameData
    {
        public string Name { get; set; }
        public string Path { get; set; }

        public string FullName
        {
            get { return string.Format("{0}\\{1}", Path, Name); }
        }
    }

    public class ObjectData<T>
    {
        public ObjectData(IntPtr address, T @object)
        {
            Address = address;
            Object = @object;
        }

        public IntPtr Address { get; set; }

        public T Object { get; set; }
    }

    public class ObjectTypeData
    {
        public string Name { get; set; }

        public ObjectData<ObjectHeaderCreatorInfo> ObjectHeaderCreatorInfo { get; set; }

        public ObjectData<ObjectHeader> ObjectHeader { get; set; }

        public ObjectData<ObjectType> ObjectType { get; set; }
    }
}

/*
REF: http://bbs.pediy.com/showthread.php?t=118325

lkd> dt _OBJECT_HEADER_AUDIT_INFO
nt!_OBJECT_HEADER_AUDIT_INFO
   +0x000 SecurityDescriptor : Ptr32 Void
   +0x004 Reserved         : Uint4B

lkd> dt _OBJECT_HEADER_PROCESS_INFO
nt!_OBJECT_HEADER_PROCESS_INFO
   +0x000 ExclusiveProcess : Ptr32 _EPROCESS
   +0x004 Reserved         : Uint4B

lkd> dt _OBJECT_HEADER_QUOTA_INFO
nt!_OBJECT_HEADER_QUOTA_INFO
   +0x000 PagedPoolCharge  : Uint4B
   +0x004 NonPagedPoolCharge : Uint4B
   +0x008 SecurityDescriptorCharge : Uint4B
   +0x00c SecurityDescriptorQuotaBlock : Ptr32 Void

lkd> dt _OBJECT_HEADER_HANDLE_INFO
nt!_OBJECT_HEADER_HANDLE_INFO
   +0x000 HandleCountDataBase : Ptr32 _OBJECT_HANDLE_COUNT_DATABASE
   +0x000 SingleEntry      : _OBJECT_HANDLE_COUNT_ENTRY

lkd> dt _OBJECT_HANDLE_COUNT_ENTRY
nt!_OBJECT_HANDLE_COUNT_ENTRY
   +0x000 Process          : Ptr32 _EPROCESS
   +0x004 HandleCount      : Pos 0, 24 Bits
   +0x004 LockCount        : Pos 24, 8 Bits

lkd> dt _OBJECT_HEADER_NAME_INFO
nt!_OBJECT_HEADER_NAME_INFO
   +0x000 Directory        : Ptr32 _OBJECT_DIRECTORY
   +0x004 Name             : _UNICODE_STRING
   +0x00c ReferenceCount   : Int4B
  
lkd> dt _OBJECT_HEADER_CREATOR_INFO
nt!_OBJECT_HEADER_CREATOR_INFO
   +0x000 TypeList         : _LIST_ENTRY
   +0x008 CreatorUniqueProcess : Ptr32 Void
   +0x00c CreatorBackTraceIndex : Uint2B
   +0x00e Reserved         : Uint2B

lkd> dt _object_header
nt!_OBJECT_HEADER
   +0x000 PointerCount     : Int4B
   +0x004 HandleCount      : Int4B
   +0x004 NextToFree       : Ptr32 Void
   +0x008 Lock             : _EX_PUSH_LOCK
   +0x00c TypeIndex        : UChar
   +0x00d TraceFlags       : UChar
   +0x00d DbgRefTrace      : Pos 0, 1 Bit
   +0x00d DbgTracePermanent : Pos 1, 1 Bit
   +0x00e InfoMask         : UChar
   +0x00f Flags            : UChar
   +0x00f NewObject        : Pos 0, 1 Bit
   +0x00f KernelObject     : Pos 1, 1 Bit
   +0x00f KernelOnlyAccess : Pos 2, 1 Bit
   +0x00f ExclusiveObject  : Pos 3, 1 Bit
   +0x00f PermanentObject  : Pos 4, 1 Bit
   +0x00f DefaultSecurityQuota : Pos 5, 1 Bit
   +0x00f SingleHandleEntry : Pos 6, 1 Bit
   +0x00f DeletedInline    : Pos 7, 1 Bit
   +0x010 ObjectCreateInfo : Ptr32 _OBJECT_CREATE_INFORMATION
   +0x010 QuotaBlockCharged : Ptr32 Void
   +0x014 SecurityDescriptor : Ptr32 Void
   +0x018 Body             : _QUAD --> #_object_type#

lkd> dt _object_type
nt!_OBJECT_TYPE
   +0x000 TypeList         : _LIST_ENTRY
   +0x008 Name             : _UNICODE_STRING
   +0x010 DefaultObject    : Ptr32 Void
   +0x014 Index            : UChar
   +0x018 TotalNumberOfObjects : Uint4B
   +0x01c TotalNumberOfHandles : Uint4B
   +0x020 HighWaterNumberOfObjects : Uint4B
   +0x024 HighWaterNumberOfHandles : Uint4B
   +0x028 TypeInfo         : _OBJECT_TYPE_INITIALIZER
   +0x080 TypeLock         : _EX_PUSH_LOCK
   +0x084 Key              : Uint4B
   +0x088 CallbackList     : _LIST_ENTRY

*/